# Copyright (c) OpenMMLab. All rights reserved.
def wizardlm_map_fn(example):
    messages = example["conversations"]
    input = ""
    conversation = []
    while messages and messages[0]["from"] == "gpt":
        # Skip the first one if it is from gpt
        messages = messages[1:]
    for msg in messages:
        if msg["from"] == "human":
            input += msg["value"]
        elif msg["from"] == "gpt":
            conversation.append({"input": input, "output": msg["value"]})
            input = ""
        else:
            raise NotImplementedError
    return {"conversation": conversation}
